package com.bestvike.stone.spring.client;

import com.bestvike.stone.core.collection.EnumerationUtils;
import com.bestvike.stone.core.reflect.ReflectionUtils;
import com.bestvike.stone.spring.context.RequestContext;
import com.bestvike.stone.spring.context.ThreadContext;
import com.bestvike.stone.spring.context.TraceContext;
import com.bestvike.stone.spring.trace.TraceHeaders;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.client.AbstractClientHttpRequest;
import org.springframework.http.client.ClientHttpRequest;
import org.springframework.http.client.ClientHttpRequestFactory;
import org.springframework.web.client.RequestCallback;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestTemplate;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLSocketFactory;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.lang.reflect.Field;
import java.net.HttpURLConnection;
import java.net.URI;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;

/**
 * 支持传递 Headers 的 RestTemplate.
 * Created by 许崇雷 on 2017-08-15.
 */
public class RestTemplateEx extends RestTemplate {
    final RestTemplateSupportedType supportedType;          //支持的类型
    boolean routeHeadersEnabled;                            //是否启用传递 Headers
    HttpHeaders ignoreHeaders;                              //忽略的 Headers

    public RestTemplateEx(RestTemplateSupportedType supportedType) {
        super(Collections.singletonList(null));
        if (supportedType == null)
            throw new NullPointerException("supportedType can not be null");
        this.supportedType = supportedType;
        //must set converters after construct
    }

    public RestTemplateEx(RestTemplateSupportedType supportedType, ClientHttpRequestFactory requestFactory) {
        this(supportedType);
        this.setRequestFactory(requestFactory);
    }

    //发起请求
    @Override
    protected <T> T doExecute(URI url, HttpMethod method, RequestCallback requestCallback, ResponseExtractor<T> responseExtractor) throws RestClientException {
        TraceContext.beginSpan();
        try {
            return super.doExecute(url, method, new RequestCallbackWrapper(requestCallback), responseExtractor);
        } finally {
            TraceContext.endSpan();
        }
    }

    //创建请求
    @Override
    protected ClientHttpRequest createRequest(URI url, HttpMethod method) throws IOException {
        ClientHttpRequest request = super.createRequest(url, method);
        return request instanceof AbstractClientHttpRequest
                ? new ClientRequestWrapper((AbstractClientHttpRequest) request)
                : request;
    }

    //静态常量
    private static class CONSTS {
        private static final Field REQUEST_ENTITY_FIELD;
        private static final Field REQUEST_FIELD;

        static {
            Class<?> clazz = null;
            try {
                clazz = Class.forName("org.springframework.web.client.RestTemplate$HttpEntityRequestCallback");
            } catch (Exception ignored) {
            }
            REQUEST_ENTITY_FIELD = ReflectionUtils.getField(clazz, "requestEntity", false);
            REQUEST_FIELD = ReflectionUtils.getField(ClientRequestWrapper.class, "request", false);
        }
    }

    //请求回调包装器
    private class RequestCallbackWrapper implements RequestCallback {
        private static final String NAME_CONNECTION = "connection";
        private static final String NAME_TOKEN = "token";
        private final RequestCallback requestCallback;

        //构造函数
        private RequestCallbackWrapper(RequestCallback requestCallback) {
            this.requestCallback = requestCallback;
        }

        //获取 HttpsEntity
        private HttpsEntity getHttpsEntity() {
            try {
                return (HttpsEntity) CONSTS.REQUEST_ENTITY_FIELD.get(this.requestCallback);
            } catch (Exception ignored) {
            }
            return null;
        }

        //获取 HttpsURLConnection
        private HttpsURLConnection getHttpsConnection(ClientHttpRequest request) {
            //拆包装
            if (request instanceof ClientRequestWrapper) {
                try {
                    request = (ClientHttpRequest) CONSTS.REQUEST_FIELD.get(request);
                } catch (Exception ignored) {
                }
            }
            //获取连接
            try {
                final HttpURLConnection connection = ReflectionUtils.getFieldValue(request, NAME_CONNECTION);
                return connection instanceof HttpsURLConnection ? (HttpsURLConnection) connection : null;
            } catch (Exception ignored) {
            }
            return null;
        }

        //准备 HttpsConnection
        private void prepareHttpsConnection(ClientHttpRequest request) {
            final HttpsEntity httpsEntity = this.getHttpsEntity();
            if (httpsEntity == null || !httpsEntity.hasSslOptions())
                return;
            final HttpsURLConnection connection = this.getHttpsConnection(request);
            if (connection == null)
                return;
            final HostnameVerifier hostnameVerifier = httpsEntity.getHostnameVerifier();
            if (hostnameVerifier != null)
                connection.setHostnameVerifier(hostnameVerifier);
            final SSLSocketFactory socketFactory = httpsEntity.getSocketFactory();
            if (socketFactory != null)
                connection.setSSLSocketFactory(socketFactory);
        }

        //设置 Headers
        private void setHeaders(ClientHttpRequest request) {
            //通用 Header
            final HttpHeaders headers = request.getHeaders();
            if (ThreadContext.exists()) {
                headers.set(TraceHeaders.TRACE_ID_NAME, TraceContext.getTraceId());
                headers.set(TraceHeaders.SPAN_ID_NAME, TraceContext.getSpanId());
                headers.set(TraceHeaders.PARENT_SPAN_ID_NAME, TraceContext.getParentSpanId());
                headers.set(NAME_TOKEN, ThreadContext.getToken());
            }
            //未启用传递 Header 功能
            if (!RestTemplateEx.this.routeHeadersEnabled)
                return;
            //原始 Header
            if (RequestContext.exists()) {
                HttpServletRequest sourceRequest = RequestContext.getRequest();
                Enumeration<String> headerNames = sourceRequest.getHeaderNames();
                while (headerNames.hasMoreElements()) {
                    String headerName = headerNames.nextElement();
                    //已有或敏感 Header 不覆盖
                    if (headers.containsKey(headerName) || RestTemplateEx.this.ignoreHeaders.containsKey(headerName))
                        continue;
                    //设置 headerName headerValues
                    List<String> headerValues = EnumerationUtils.toList(sourceRequest.getHeaders(headerName));
                    headers.put(headerName, headerValues);
                }
            }
        }

        //执行回调
        @Override
        public void doWithRequest(ClientHttpRequest request) throws IOException {
            if (this.requestCallback != null) {
                this.prepareHttpsConnection(request);
                this.requestCallback.doWithRequest(request);
            }
            this.setHeaders(request);
        }
    }
}
